from pocketflow import Node
from tools.vision import extract_text_from_image
from pathlib import Path
import os
from utils.call_llm import call_llm
from utils.content_utils import extract_html_blocks, extract_json_blocks
from utils.prompt import GEN_HTML_PROMPT, ANNLYSIS_PROMPT

class ExtractTextNode(Node):
    """Node for extracting text from images using Vision API"""
    
    def prep(self, shared):
        root_dir = Path(__file__).parent
        img_path= ''
        if shared.get("img_path", ''):
            img_path = root_dir/shared["img_path"]
            inputs = {
                        "user_input": shared["question"],
                        "img_path": str(img_path),
                        "extraction_prompt": "Extract all text from this document, preserving formatting and layout."
                    }
        else:
            inputs = {
                        "user_input": shared["question"],
                        "img_path": None,
                        "extraction_prompt": None
                    }
        print(f"Process images: {inputs}")
        return inputs
        
    def exec(self, prep_res):
        user_input = prep_res["user_input"]
        img_path = prep_res["img_path"]
        extraction_prompt = prep_res["extraction_prompt"]
        text = None
        # 检查图片文件是否存在
        if img_path and os.path.exists(img_path):
            # 使用Pillow打开图片文件
            from PIL import Image
            img = Image.open(img_path)
            
            text = extract_text_from_image(img, extraction_prompt)
        # print(user_input)
        print("Extract text from image: ",text)
        return user_input, text
        
    def post(self, shared, prep_res, exec_res):
        user_input, text = exec_res
        shared["user_input"] = user_input
        if text:
            shared["user_input"] = user_input + ", " + text
            
        return 'default'
    
class AnalysisContentNode(Node):
    """Node for analyzing content using LLM API"""
    
    def prep(self, shared):
        return shared["user_input"]
        
        
    def exec(self, prep_res):
        print("Analysis content: ", prep_res)
        user_input = prep_res
        # analysis_prompt.replace("{{userInput}}", user_input)

        content = call_llm(ANNLYSIS_PROMPT, user_input)
        print("Analysis llm result: ", content)
        analysis_result = extract_json_blocks(content)
        if analysis_result:
            return analysis_result[0]
        return ''
        
    def post(self, shared, prep_res, exec_res):
        # print("Analysis result: ", exec_res)
        shared["analysis_result"] = exec_res
        return "default"


class GenerateHtmlNode(Node):
    """Node for combining and formatting extracted text"""
    
    def prep(self, shared):
        return shared["analysis_result"]
        
    def exec(self, prep_res):
        # 使用LLM API生成HTML内容
        html_content = call_llm(GEN_HTML_PROMPT, prep_res)
        print("Html content from LLM: ", html_content)
        if html_content.startswith("<!DOCTYPE html>"):
            return html_content
        extract_result = extract_html_blocks(html_content)
        if extract_result:
            return extract_result[0]
        return ''
        
    def post(self, shared, prep_res, exec_res):
        print("Html content: ",exec_res)
        # 保存到static/gen目录下面
        import time
        import hashlib
        
        # 生成时间戳+随机MD5的文件名
        timestamp = int(time.time())
        random_str = hashlib.md5(os.urandom(32)).hexdigest()[:8]
        file_name = f"{timestamp}_{random_str}.html"
        
        # 确保目录存在
        output_dir = Path(__file__).parent / "static" / "gen"
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # 写入HTML文件
        output_path = output_dir / file_name
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(exec_res)
        print("Html content path: ",output_path)
        shared["final_text"] = str("static/gen/" + file_name)
        return "default"
